package com.ti.demo.configuration;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SafeGuardAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.VectorStoreChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;

import java.util.Arrays;

/**
 * @description: com.ti.demo.configuration
 * @Author Triagen
 * @Date 2025/7/18 11:12
 */
@Configuration
public class ChatClientConfiguration {
    private static final String DEFAULT_SENSITIVE_WORDS = "澳门赌博,天安门事件,文化大革命,苍老师,大麻,鸦片,冰毒";

    @Autowired
    private ChatMemory chatMemory;

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private ChatClient.Builder chatClientBuilder;

    @Bean("ollamaChatClient")
    public ChatClient ollamaChatClient(OllamaChatModel ollamaChatModel) {
        return ChatClient.builder(ollamaChatModel).defaultSystem("You are a helpful image recognition assistant.").build();
    }


    /**
     * a manual ChatClient example,for better understanding of api, model and chatClient
     * generally speaking, api includes the connection info, model includes the model choice, and chatClient is the assembly of both
     *
     * @param openAiChatModel
     * @return
     */
    @Bean("openaiChatClient")
    public ChatClient openaiChatClient(OpenAiChatModel openAiChatModel) {
        return ChatClient.builder(openAiChatModel)
                .defaultAdvisors(
                        SafeGuardAdvisor.builder().order(Ordered.HIGHEST_PRECEDENCE + 999).sensitiveWords(Arrays.asList(DEFAULT_SENSITIVE_WORDS.split(","))).failureResponse("您的问题太敏感，请换个问题").build(),
                        SimpleLoggerAdvisor.builder().build(),
                        MessageChatMemoryAdvisor.builder(chatMemory).build())
                .defaultSystem("You are a good joker.").build();
    }


    @Bean("poetClient")
    public ChatClient poetClient(OpenAiChatModel openAiChatModel) {
        return ChatClient.builder(openAiChatModel)
                .defaultAdvisors(
                        SafeGuardAdvisor.builder().order(Ordered.HIGHEST_PRECEDENCE + 999).sensitiveWords(Arrays.asList(DEFAULT_SENSITIVE_WORDS.split(","))).failureResponse("您的问题太敏感，请换个问题").build(),
                        SimpleLoggerAdvisor.builder().build(),
                        VectorStoreChatMemoryAdvisor.builder(vectorStore).defaultTopK(10).build())
                .defaultSystem("You are a poet and good at Chinese poem.").build();
    }

    @Bean("qaClient")
    public ChatClient qaClient(OpenAiChatModel openAiChatModel) {
        return ChatClient.builder(openAiChatModel)
                .defaultAdvisors(
                        SafeGuardAdvisor.builder().order(Ordered.HIGHEST_PRECEDENCE + 999).sensitiveWords(Arrays.asList(DEFAULT_SENSITIVE_WORDS.split(","))).failureResponse("您的问题太敏感，请换个问题").build(),
                        SimpleLoggerAdvisor.builder().build(),
                        QuestionAnswerAdvisor.builder(vectorStore).searchRequest(SearchRequest.builder().similarityThreshold(0.3).topK(1).build()).build())
                .defaultSystem("You are good at answer questions.").build();
    }

    @Bean("raClient")
    public ChatClient raClient(OpenAiChatModel openAiChatModel) {
        return ChatClient.builder(openAiChatModel)
                .defaultAdvisors(
                        SafeGuardAdvisor.builder().order(Ordered.HIGHEST_PRECEDENCE + 999).sensitiveWords(Arrays.asList(DEFAULT_SENSITIVE_WORDS.split(","))).failureResponse("您的问题太敏感，请换个问题").build(),
                        SimpleLoggerAdvisor.builder().build(),
                        RetrievalAugmentationAdvisor.builder()
                                .queryTransformers(TranslationQueryTransformer.builder().chatClientBuilder(chatClientBuilder).targetLanguage("Chinese").build())
                                .queryExpander(MultiQueryExpander.builder().chatClientBuilder(chatClientBuilder).numberOfQueries(5).build())
                                .documentRetriever(VectorStoreDocumentRetriever.builder().similarityThreshold(0.5).topK(4).vectorStore(vectorStore).build())
                                .build())
                .defaultSystem("You are good at knowledge retrieval.").build();
    }


}
